from BinaryTree import *


class ArrayBinaryTree(BinaryTree):
    """数组表示的二叉树"""

    class Position(BinaryTree.Position):
        """每个位置是一个节点"""

        def __init__(self, container, index):
            self.container = container
            self.index = index
            self._element = container._data[index]
            self._parent = self._parent(index)
            self._left = self._left(container, index)
            self._right = self._right(container, index)

        def element(self):
            return self._element

        def _parent(self,index):
            if index == 0:
                return None
            return (index - 1) // 2

        def _left(self, container, index):
            left = 2 * index + 1
            if left >= len(container._data):
                return None
            elif container._data[left] is None:
                return None
            else:
                return left

        def _right(self, container, index):
            right = 2 * index + 2
            if right >= len(container._data):
                return None
            elif container._data[right] is None:
                return None
            else:
                return right

        def __eq__(self, other):
            return type(other) is type(self) and other.index == self.index

    # --------------------------构造函数---------------------------
    def __init__(self):
        self._data = [None]
        self._root = None
        self._size = 0

    # --------------------------公共方法-----------------------------
    def __len__(self):
        return self._size

    def __str__(self):
        msg = str(self._data)
        return msg

    def validate(self, p):
        """判断位置p是否合法"""
        if not isinstance(p, self.Position):
            raise TypeError('p不是一个位置对象')
        if p.container is not self:
            raise ValueError('p不属于该树')
        index = p.index
        if p._parent is index:
            raise ValueError('p is no longer valid')
        if index < 0 or index > len(self._data):
            raise ValueError('p位置超出范围')
        if self._data[index] is None:
            raise ValueError('p位置不存在')
        return index

    def make_position(self, index):
        return self.Position(self, index) if index is not None else None

    def index(self, p):
        """层编号函数,顺便判断是否合法节点"""
        return p.index

    def root(self):
        return self._root

    def parent(self, p):
        self.validate(p)
        if p is not self.root():
            return self.make_position(p._parent)

    def left(self, p):
        self.validate(p)
        return self.make_position(p._left)

    def right(self, p):
        self.validate(p)
        return self.make_position(p._right)

    def num_children(self,p):
        self.validate(p)
        left = p._left
        right = p._right
        count = 0
        if 0 <= left < len(self._data):
            if self._data[left] is not None:
                count += 1
        if 0 <= right < len(self._data):
            if self._data[right] is not None:
                count += 1
        return count

    def add_root(self, e):
        if self._root is not None:
            raise ValueError('root已存在')
        self._data[0] = e
        self._size += 1
        self._root = self.make_position(0)
        return self._root

    def add_left(self, p, e):
        self.validate(p)
        left = p._left
        if left is None:
            left = p.index * 2 + 1
        if left >= len(self._data):
            self.resize(left + 1)
        if self._data[left] is not None:
            raise ValueError('left 已存在')
        p._left = left
        self._data[left] = e
        self._size += 1
        return self.make_position(left)

    def add_right(self, p, e):
        self.validate(p)
        right = p._right
        if right is None:
            right = p.index * 2 + 2
        if right >= len(self._data):
            self.resize(right + 1)
        if self._data[right] is not None:
            raise ValueError('right 已存在')
        p._right = right
        self._data[right] = e
        self._size += 1
        return self.make_position(right)

    def resize(self, length):
        tem_data = [None] * (2 * len(self._data))
        while len(tem_data) < length:
            tem_data = [None] * (2 * len(tem_data))
        for index in range(len(self._data)):
            tem_data[index] = self._data[index]
        self._data = tem_data


if __name__ == '__main__':
    T = ArrayBinaryTree()
    T.add_root(9)
    T.add_left(T.root(),2)
    T.add_right(T.root(),3)
    print(T)
    for e in T:
        print(e)
